from __future__ import annotations

import os
import uuid

from typing import Dict, List, Union, Optional
from common.typing_utils import numeric

from osgeo import gdal

from common.config import DEFAULT_CACHE_DIR
from common.logger import logger
from common.io_utils.dirs import mkdir_p

from common.geometry.bbox import BoundingBox
from common.geoimage.gdal_sys import SceneMeta, read_meta_from_ds
from common.geoimage.utils import (
    polyfill_resample_alg_for_gdal_vrt,
    SUPPORTED_WARP_RESAMPLE_ALGS,
)

TMP_DIR = os.path.join(DEFAULT_CACHE_DIR, "common", "geoimage", "virtual_raster")
mkdir_p(TMP_DIR)


class VRTScene(object):
    def __init__(self):
        self._vrt_file_path: Optional[str] = None
        self._meta: Optional[SceneMeta] = None

    @classmethod
    def from_vrt_file_path(cls: VRTScene, vrt_file_path: str) -> VRTScene:
        """Create a VRTScene from a VRT file path.

        Parameters
        ----------
        vrt_file_path : str
            File path to the gdal Virtual Raster (VRT) file.

        Returns
        -------
        VRTScene
            The created `VRTScene` object.
        """
        vrt_scene = cls()
        vrt_scene._vrt_file_path = vrt_file_path

        ds = gdal.Open(vrt_file_path)
        vrt_scene._meta = read_meta_from_ds(ds, read_value_interpretation=True)
        ds = None
        return vrt_scene

    @classmethod
    def from_mosaic_files(
        cls: VRTScene,
        lst_files: List[str],
        resample_alg: str = "near",
        src_nodata: Union[Optional[numeric], List[Optional[numeric]]] = None,
        dest_extent: Optional[BoundingBox] = None,
        dest_res: Optional[numeric] = None,
        dest_nodata: Union[Optional[numeric], List[Optional[numeric]]] = None,
    ) -> VRTScene:
        """Create a VRTScene by spatially mosaic a list of image files,
        assuming they are under the same coordinate system.

        Parameters
        ----------
        lst_files : List[str]
            A list of file paths to gdal supported images to be mosaiced.
        resample_alg : str, optional
            Resampling algorithm to be used, could be one of
            "nearest", "bilinear", "cubic", "cubicspline", "lanczos", "average", and "mode",
            by default "nearest".
        src_nodata : Optional[Union[numeric, List[numeric]]], optional
            Source data's nodata value(s),
            could be a single value for all bands, or a list of values with each value for one band,
            by default None to let `gdal` determine the nodata value(s).
        dest_extent : Optional[BoundingBox], optional
            Destination extent as under the result spatial reference system, by default None
        dest_res: Optional[numeric], optional
            Destination resolution, by default None
        dest_nodata : Optional[Union[numeric, List[numeric]]], optional
            Destination data's nodata value(s),
            could be a single value for all bands, or a list of values with each value for one band,
            by default None to use the source data's nodata value(s).

        Returns
        -------
        VRTScene
            Mosaiced `VRTScene` object.

        Raises
        ------
        err
            Raised when the call to `gdal.BuildVRT` failed.

        Notes
        -----
        This function expects all given raster files or datasets are in the same spatial reference system,
        and expects they have the same number of bands; if the `dest_extent` and `dest_res` are provided,
        those are expected to be compatible with the sources' coordinate system.
        """
        tmp_uuid = str(uuid.uuid4())
        fpath_vrt = os.path.join(TMP_DIR, f"mosaiced_{tmp_uuid}.vrt")

        resample_alg = polyfill_resample_alg_for_gdal_vrt(resample_alg)
        if dest_extent is not None:
            output_bounds = dest_extent.to_gdal_target_extent()
        else:
            output_bounds = None

        try:
            ds = gdal.BuildVRT(
                fpath_vrt,
                lst_files,
                options=gdal.BuildVRTOptions(
                    resolution="highest" if dest_res is None else "user",
                    resampleAlg=resample_alg,
                    separate=False,
                    allowProjectionDifference=False,
                    srcNodata=src_nodata,
                    VRTNodata=dest_nodata,
                    outputSRS=None,
                    outputBounds=output_bounds,
                    xRes=dest_res,
                    yRes=dest_res,
                ),
            )
            del ds
        except Exception as err:
            ds = None  # noqa F841
            logger.error(f"Failed to build VRT due to {str(err)}")
            raise err
        else:
            vrt_scene = cls.from_vrt_file_path(fpath_vrt)
            return vrt_scene

    @classmethod
    def from_stack_files(
        cls: VRTScene,
        lst_files: List[str],
        resample_alg: str = "near",
        src_nodata: Union[Optional[numeric], List[Optional[numeric]]] = None,
        dest_extent: Optional[BoundingBox] = None,
        dest_res: Optional[numeric] = None,
        dest_nodata: Union[Optional[numeric], List[Optional[numeric]]] = None,
    ) -> VRTScene:
        """Create a VRTScene by stacking a list of image files along the band's dimension.

        Parameters
        ----------
        lst_files : List[str]
            A list of file paths to gdal supported images to be mosaiced.
        resample_alg : str, optional
            Resampling algorithm to be used, could be one of
            "nearest", "bilinear", "cubic", "cubicspline", "lanczos", "average", and "mode",
            by default "nearest".
        src_nodata : Optional[Union[numeric, List[numeric]]], optional
            Source data's nodata value(s),
            could be a single value for all bands, or a list of values with each value for one band,
            by default None to let `gdal` determine the nodata value(s).
        dest_extent : Optional[BoundingBox], optional
            Destination extent as under the result spatial reference system, by default None
        dest_res: Optional[numeric], optional
            Destination resolution, by default None
        dest_nodata : Optional[Union[numeric, List[numeric]]], optional
            Destination data's nodata value(s),
            could be a single value for all bands, or a list of values with each value for one band,
            by default None to use the source data's nodata value(s).

        Returns
        -------
        VRTScene
            Stacked `VRTScene` object.

        Raises
        ------
        err
            Raised when the call to `gdal.BuildVRT` failed.

        Notes
        -----
        This function expects all given raster files are in the same spatial reference system,
        and expects them to be roughly within the same spatial extent.
        """
        # TODO, same for the from_mosaic_files
        # 1. maybe a better strategy of making a hash from assets + further operations
        #    or maybe just a in-memory only VRT
        # 2. could optionally ask for more info
        #    so can create the VRT by writing our own XML instead of go through BuildVRT to reduce IO
        #    especially if later we adopt the STAC's raster extension then we will have all info needed
        tmp_uuid = str(uuid.uuid4())
        fpath_vrt = os.path.join(TMP_DIR, f"stacked_{tmp_uuid}.vrt")

        resample_alg = polyfill_resample_alg_for_gdal_vrt(resample_alg)
        if dest_extent is not None:
            output_bounds = dest_extent.to_gdal_target_extent()
        else:
            output_bounds = None

        if isinstance(src_nodata, list):
            src_nodata = " ".join([str(v) for v in src_nodata])

        if isinstance(dest_nodata, list):
            dest_nodata = " ".join([str(v) for v in dest_nodata])

        try:
            ds = gdal.BuildVRT(
                fpath_vrt,
                lst_files,
                options=gdal.BuildVRTOptions(
                    resolution="highest" if dest_res is None else "user",
                    resampleAlg=resample_alg,
                    separate=True,
                    allowProjectionDifference=False,
                    srcNodata=src_nodata,
                    VRTNodata=dest_nodata,
                    outputSRS=None,
                    outputBounds=output_bounds,
                    xRes=dest_res,
                    yRes=dest_res,
                ),
            )
            del ds
        except Exception as err:
            ds = None  # noqa F841
            logger.error(f"Failed to build VRT due to {str(err)}")
            raise err
        else:
            vrt_scene = cls.from_vrt_file_path(fpath_vrt)
            return vrt_scene

    @classmethod
    def from_assets(
        cls: VRTScene,
        asset_paths: Dict[str, str],
        resample_alg: str = "near",
        src_nodata: Union[Optional[numeric], List[Optional[numeric]]] = None,
        dest_extent: Optional[BoundingBox] = None,
        dest_res: Optional[numeric] = None,
        dest_nodata: Union[Optional[numeric], List[Optional[numeric]]] = 0,
    ) -> VRTScene:
        """Create a VRTScene from a dictionary of assets.

        Parameters
        ----------
        asset_paths : Dict[str, str]
            A dictionary, with the key as the asset name, and value as the file path of the asset.
        resample_alg : str, optional
            Resampling algorithm to be used, could be one of
            "nearest", "bilinear", "cubic", "cubicspline", "lanczos", "average", and "mode",
            by default "nearest".
        src_nodata : Optional[Union[numeric, List[numeric]]], optional
            Source data's nodata value(s),
            could be a single value for all bands, or a list of values with each value for one band,
            by default None to let `gdal` determine the nodata value(s).
        dest_extent : Optional[BoundingBox], optional
            Destination extent as under the result spatial reference system, by default None
        dest_res: Optional[numeric], optional
            Destination resolution, by default None
        dest_nodata : Optional[Union[numeric, List[numeric]]], optional
            Destination data's nodata value(s),
            could be a single value for all bands, or a list of values with each value for one band,
            by default None to use the source data's nodata value(s).

        Returns
        -------
        VRTScene
            Stacked `VRTScene` object.
        """
        # TODO support specifying assets and bands within assets,
        # this might need to change the API and affecting downstream code
        # this would require `from_stack_files` to support a `band_list_in_files` argument
        # TODO more precise control of the VRT generation
        lst_files = [fpath for fpath in asset_paths.values()]
        return cls.from_stack_files(
            lst_files,
            resample_alg=resample_alg,
            src_nodata=src_nodata,
            dest_extent=dest_extent,
            dest_res=dest_res,
            dest_nodata=dest_nodata,
        )

    def warp(
        self,
        dest_srs: Optional[str] = None,
        dest_extent: Optional[BoundingBox] = None,
        dest_x_res: Optional[numeric] = None,
        dest_y_res: Optional[numeric] = None,
        dest_nodata: Union[Optional[numeric], List[Optional[numeric]]] = 0,
        resample_alg: str = "near",
        target_aligned_pixels: bool = False,
    ) -> VRTScene:
        """Warp a VRTScene to a new VRTScene with different spatial reference system, or extent, or resolution.

        Parameters
        ----------
        dest_srs : Optional[str], optional
            _description_, by default None
        dest_extent : Optional[BoundingBox], optional
            Destination extent as under the result spatial reference system, by default None
        dest_x_res: Optional[numeric], optional
            Destination resolution in the x direction, by default None
        dest_y_res: Optional[numeric], optional
            Destination resolution in the y direction, by default None
        dest_nodata : Optional[Union[numeric, List[numeric]]], optional
            Destination data's nodata value(s),
            could be a single value for all bands, or a list of values with each value for one band,
            by default None to use the source data's nodata value(s).
        resample_alg : str, optional
            Resampling algorithm to be used, could be one of
            "nearest", "bilinear", "cubic", "cubicspline", "lanczos", "average", and "mode",
            by default "nearest".
        target_aligned_pixels : bool, optional
            Whether to align target pixels, by default False not to align.
            NOTE: This is instruct gdal to use its internal way of aligning pixels
            but users frequently find it not behaving as they expect.
            Thus, please use this option with caution and only when you know what you are doing.
            In general, We opt to align pixels by manipulating and setting
            the right combination of destination extent and destination resolutions.

        Returns
        -------
        VRTScene
            Mosaiced `VRTScene` object.

        Raises
        ------
        err
            Raised when the call to `gdal.Warp` failed.
        """
        if resample_alg not in SUPPORTED_WARP_RESAMPLE_ALGS:
            raise ValueError(
                f"{resample_alg} is not one of {SUPPORTED_WARP_RESAMPLE_ALGS}."
            )

        # TODO
        # 1. support more options to allow more precise control
        # 2. write our own VRT instead of calling gdal to reduce IO
        if dest_srs is None:
            dest_srs = self._meta.prj_epsg
        dest_srs_epsg_no = int(dest_srs.split(":")[-1])
        src_srs_epsg_no = int(self._meta.prj_epsg.split(":")[-1])
        if (src_srs_epsg_no == dest_srs_epsg_no) and (not target_aligned_pixels):
            if (dest_x_res is None) and (dest_y_res is None):
                return self
            elif dest_x_res == abs(self._meta.x_res) and dest_y_res == abs(
                self._meta.y_res
            ):
                return self

        if dest_extent is not None:
            output_bounds = dest_extent.to_gdal_target_extent()
        else:
            output_bounds = None

        tmp_uuid = str(uuid.uuid4())
        vrt_file_path = os.path.join(TMP_DIR, "warped_{}.vrt".format(tmp_uuid))
        ds = gdal.Warp(
            vrt_file_path,
            self._vrt_file_path,
            options=gdal.WarpOptions(
                dstSRS=dest_srs,
                xRes=dest_x_res,
                yRes=dest_y_res,
                outputBounds=output_bounds,
                resampleAlg=resample_alg,
                targetAlignedPixels=target_aligned_pixels,
                dstNodata=dest_nodata,
            ),
        )
        del ds
        return VRTScene.from_vrt_file_path(vrt_file_path)


class VRTSceneCollection(object):
    """An in-memory raster dataset

    Attributes
    ----------
    lst_scenes: List[VRTScene]
        A list of VRTScene objects in the collection.
    """

    # NOTE this is the currently recommended way of type annotating **an instance variable**
    # as noted by PEP-0526: https://www.python.org/dev/peps/pep-0526/#class-and-instance-variable-annotations
    # and supported by the BDFL himself: https://gitter.im/python/typing?at=595f9f5b2723db8d5ee9c892
    # this is kind of confusing as it is with the same syntax as the traditional way of defining a class variable
    # TODO, add a "Developer Note" section in our `skywalker-docs` and collect more notes there
    lst_scenes: List[VRTScene] = []

    def __init__(self, lst_scenes):
        self.lst_scenes = lst_scenes

    def warp_to_common_srs(
        self,
        dest_srs: Optional[str],
        dest_x_res=None,
        dest_y_res=None,
        dest_nodata: Union[Optional[numeric], List[Optional[numeric]]] = 0,
        resample_alg: str = "near",
        target_aligned_pixels=False,
    ) -> VRTScene:
        """Warp all VRTScene objects in the collection to a common spatial reference system.

        Parameters
        ----------
        dest_srs : Optional[str], optional
            _description_, by default None
        dest_x_res: Optional[numeric], optional
            Destination resolution in the x direction, by default None
        dest_y_res: Optional[numeric], optional
            Destination resolution in the y direction, by default None
        dest_nodata : Optional[Union[numeric, List[numeric]]], optional
            Destination data's nodata value(s),
            could be a single value for all bands, or a list of values with each value for one band,
            by default None to use the source data's nodata value(s).
        resample_alg : str, optional
            Resampling algorithm to be used, could be one of
            "nearest", "bilinear", "cubic", "cubicspline", "lanczos", "average", and "mode",
            by default "nearest".
        target_aligned_pixels : bool, optional
            Whether to align target pixels, by default False not to align.
            NOTE: This is instruct gdal to use its internal way of aligning pixels
            but users frequently find it not behaving as they expect.
            Thus, please use this option with caution and only when you know what you are doing.
            In general, We opt to align pixels by manipulating and setting
            the right combination of destination extent and destination resolutions.

        Returns
        -------
        VRTScene
            _description_
        """
        lst_scenes_warped = []
        for scene in self.lst_scenes:
            lst_scenes_warped.append(
                scene.warp(
                    dest_srs,
                    dest_x_res=dest_x_res,
                    dest_y_res=dest_y_res,
                    dest_nodata=dest_nodata,
                    resample_alg=resample_alg,
                    target_aligned_pixels=target_aligned_pixels,
                )
            )
        return VRTSceneCollection(lst_scenes_warped)

    def mosaic_vrt(
        self,
        # vrt_nodata: Union[Optional[numeric], List[Optional[numeric]]] = 0,
    ) -> VRTSceneCollection:
        """Spatially mosaic the VRTScene objects in the collection into a new VRTScene object."""
        # TODO
        # 1. support more precise control, e.g source and destination nodata values
        # turn off the vrt_nodata since we don't actually need it right now
        # and VRT has some special requirements for specifying different nodata values
        # for different bands when the source is a set of VRTs
        # 2. if time allowed, write our own VRT generation code instead of calling gdal.BuildVRT
        tmp_uuid = str(uuid.uuid4())
        vrt_file_path = os.path.join(TMP_DIR, "mosaic_{}.vrt".format(tmp_uuid))

        lst_ds = [vrt_scene._vrt_file_path for vrt_scene in self.lst_scenes]

        # if isinstance(vrt_nodata, list):
        #     vrt_nodata = " ".join([str(v) for v in vrt_nodata])
        #     vrt_nodata = '"{}"'.format(vrt_nodata)

        ds = gdal.BuildVRT(
            vrt_file_path,
            lst_ds,
            options=gdal.BuildVRTOptions(
                resolution="highest",
                # VRTNodata=vrt_nodata,
            ),
        )
        del ds
        return VRTScene.from_vrt_file_path(vrt_file_path)
